class Solution {
    private int[] p;

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        p = new int[n];
        int[] size = new int[n];
        for (int i = 0; i < n; ++i) {
            p[i] = i;
            size[i] = 1;
        }
        boolean[] clean = new boolean[n];
        Arrays.fill(clean, true);
        for (int i : initial) {
            clean[i] = false;
        }
        for (int i = 0; i < n; ++i) {
            if (!clean[i]) {
                continue;
            }
            for (int j = i + 1; j < n; ++j) {
                if (!clean[j]) {
                    continue;
                }
                if (graph[i][j] == 1) {
                    int pa = find(i), pb = find(j);
                    if (pa == pb) {
                        continue;
                    }
                    p[pa] = pb;
                    size[pb] += size[pa];
                }
            }
        }
        int[] cnt = new int[n];
        Map<Integer, Set<Integer>> mp = new HashMap<>();
        for (int i : initial) {
            Set<Integer> s = new HashSet<>();
            for (int j = 0; j < n; ++j) {
                if (!clean[j]) {
                    continue;
                }
                if (graph[i][j] == 1) {
                    s.add(find(j));
                }
            }
            for (int e : s) {
                cnt[e] += 1;
            }
            mp.put(i, s);
        }
        int mx = -1;
        int res = 0;
        for (Map.Entry<Integer, Set<Integer>> entry : mp.entrySet()) {
            int i = entry.getKey();
            int t = 0;
            for (int e : entry.getValue()) {
                if (cnt[e] == 1) {
                    t += size[e];
                }
            }
            if (mx < t || (mx == t && i < res)) {
                mx = t;
                res = i;
            }
        }
        return res;
    }

    private int find(int x) {
        if (p[x] != x) {
            p[x] = find(p[x]);
        }
        return p[x];
    }
}